// This file is part of SVO - Semi-direct Visual Odometry.
//
// Copyright (C) 2014 Christian Forster <forster at ifi dot uzh dot ch>
// (Robotics and Perception Group, University of Zurich, Switzerland).
//
// SVO is free software: you can redistribute it and/or modify it under the
// terms of the GNU General Public License as published by the Free Software
// Foundation, either version 3 of the License, or any later version.
//
// SVO is distributed in the hope that it will be useful, but WITHOUT ANY
// WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
// FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#include <svo_ros/visualizer.h>
#include <svo/frame_handler_mono.h>
#include <svo/frame.h>
#include <svo/point.h>
#include <svo/depth_filter.h>
#include <svo/map.h>
#include <svo/feature.h>
#include <svo_msgs/Info.h>
#include <geometry_msgs/PoseWithCovarianceStamped.h>
#include <sensor_msgs/image_encodings.h>
#include <cv_bridge/cv_bridge.h>
#include <tf/tf.h>
#include <iostream>
#include <fstream>
#include <ros/package.h>
#include <vikit/timer.h>
#include <vikit/output_helper.h>
#include <vikit/params_helper.h>
#include <deque>
#include <algorithm>
#include <svo_msgs/DenseInput.h>
#include <ros/ros.h>
#include <visualization_msgs/Marker.h>

namespace svo {

        Visualizer::
                Visualizer() :
                        pnh_("~"),
                        trace_id_(0),
                        img_pub_level_(vk::getParam<int>("svo/publish_img_pyr_level", 0)),
                        img_pub_nth_(vk::getParam<int>("svo/publish_every_nth_img", 1)),
                        dense_pub_nth_(vk::getParam<int>("svo/publish_every_nth_dense_input", 1)),
                        publish_world_in_cam_frame_(vk::getParam<bool>("svo/publish_world_in_cam_frame", true)),
                        publish_map_every_frame_(vk::getParam<bool>("svo/publish_map_every_frame", false)),
                        publish_points_display_time_(vk::getParam<double>("svo/publish_point_display_time", 0)),
                        T_world_from_vision_(Matrix3d::Identity(), Vector3d::Zero())
        {
                // Init ROS Marker Publishers
                pub_frames_ = pnh_.advertise<visualization_msgs::Marker>("keyframes", 10);
                pub_points_ = pnh_.advertise<visualization_msgs::Marker>("points", 1000);
                pub_pose_ = pnh_.advertise<geometry_msgs::PoseWithCovarianceStamped>("pose",10);
                pub_info_ = pnh_.advertise<svo_msgs::Info>("info", 10);
                pub_dense_ = pnh_.advertise<svo_msgs::DenseInput>("dense_input",10);

                // create video publisher
                image_transport::ImageTransport it(pnh_);
                pub_images_ = it.advertise("image", 10);
        }

        void Visualizer::publishMinimal(
                        const cv::Mat& img,
                        const FramePtr& frame,
                        const FrameHandlerMono& slam,
                        const double timestamp)
        {
                ++trace_id_;
                std_msgs::Header header_msg;
                header_msg.frame_id = "/cam";
                header_msg.seq = trace_id_;
                header_msg.stamp = ros::Time(timestamp);

                // publish info msg.
                if(pub_info_.getNumSubscribers() > 0)
                {
                        svo_msgs::Info msg_info;
                        msg_info.header = header_msg;
                        msg_info.processing_time = slam.lastProcessingTime();
                        msg_info.keyframes.resize(slam.map().keyframes_.size());
                        for(list<FramePtr>::const_iterator it=slam.map().keyframes_.begin(); it!=slam.map().keyframes_.end(); ++it)
                                msg_info.keyframes.push_back((*it)->id_);
                        msg_info.stage = static_cast<int>(slam.stage());
                        msg_info.tracking_quality = static_cast<int>(slam.trackingQuality());
                        if(frame != NULL)
                                msg_info.num_matches = slam.lastNumObservations();
                        else
                                msg_info.num_matches = 0;
                        pub_info_.publish(msg_info);    // #pub 1
                }

                if(frame == NULL)
                {
                        if(pub_images_.getNumSubscribers() > 0 && slam.stage() == FrameHandlerBase::STAGE_PAUSED)
                        {
                                // Display image when slam is not running.
                                cv_bridge::CvImage img_msg;
                                img_msg.header.stamp = ros::Time::now();
                                img_msg.header.frame_id = "/image";
                                img_msg.image = img;
                                img_msg.encoding = sensor_msgs::image_encodings::MONO8;
                                pub_images_.publish(img_msg.toImageMsg());
                        }
                        return;
                }

                // Publish pyramid-image every nth-frame.
                if(img_pub_nth_ > 0 && trace_id_%img_pub_nth_ == 0 && pub_images_.getNumSubscribers() > 0)
                {
                        const int scale = (1<<img_pub_level_);
                        cv::Mat img_rgb(frame->img_pyr_[img_pub_level_].size(), CV_8UC3);
                        cv::cvtColor(frame->img_pyr_[img_pub_level_], img_rgb, CV_GRAY2RGB);

                        if(slam.stage() == FrameHandlerBase::STAGE_SECOND_FRAME)
                        {
                                // During initialization, draw lines.
                                const vector<cv::Point2f>& px_ref(slam.initFeatureTrackRefPx());
                                const vector<cv::Point2f>& px_cur(slam.initFeatureTrackCurPx());
                                for(vector<cv::Point2f>::const_iterator it_ref=px_ref.begin(), it_cur=px_cur.begin();
                                                it_ref != px_ref.end(); ++it_ref, ++it_cur)
                                        cv::line(img_rgb,
                                                        cv::Point2f(it_cur->x/scale, it_cur->y/scale),
                                                        cv::Point2f(it_ref->x/scale, it_ref->y/scale), cv::Scalar(0,255,0), 2);
                        }

                        if(img_pub_level_ == 0)
                        {
                                for(Features::iterator it=frame->fts_.begin(); it!=frame->fts_.end(); ++it)
                                {
                                        if((*it)->type == Feature::EDGELET)
                                                cv::line(img_rgb,
                                                                cv::Point2f((*it)->px[0]+3*(*it)->grad[1], (*it)->px[1]-3*(*it)->grad[0]),
                                                                cv::Point2f((*it)->px[0]-3*(*it)->grad[1], (*it)->px[1]+3*(*it)->grad[0]),
                                                                cv::Scalar(255,0,255), 2);
                                        else
                                                cv::rectangle(img_rgb,
                                                                cv::Point2f((*it)->px[0]-2, (*it)->px[1]-2),
                                                                cv::Point2f((*it)->px[0]+2, (*it)->px[1]+2),
                                                                cv::Scalar(0,255,0), CV_FILLED);
                                }
                        }
                        if(img_pub_level_ == 1)
                                for(Features::iterator it=frame->fts_.begin(); it!=frame->fts_.end(); ++it)
                                        cv::rectangle(img_rgb,
                                                        cv::Point2f((*it)->px[0]/scale-1, (*it)->px[1]/scale-1),
                                                        cv::Point2f((*it)->px[0]/scale+1, (*it)->px[1]/scale+1),
                                                        cv::Scalar(0,255,0), CV_FILLED);
                        else
                                for(Features::iterator it=frame->fts_.begin(); it!=frame->fts_.end(); ++it)
                                        cv::rectangle(img_rgb,
                                                        cv::Point2f((*it)->px[0]/scale, (*it)->px[1]/scale),
                                                        cv::Point2f((*it)->px[0]/scale, (*it)->px[1]/scale),
                                                        cv::Scalar(0,255,0), CV_FILLED);

                        cv_bridge::CvImage img_msg;
                        img_msg.header = header_msg;
                        img_msg.image = img_rgb;
                        img_msg.encoding = sensor_msgs::image_encodings::BGR8;
                        pub_images_.publish(img_msg.toImageMsg());  // #pub 2
                }

                if(pub_pose_.getNumSubscribers() > 0 && slam.stage() == FrameHandlerBase::STAGE_DEFAULT_FRAME)
                {
                        Quaterniond q;
                        Vector3d p;
                        Matrix<double,6,6> Cov;
                        if(publish_world_in_cam_frame_)
                        {
                                // publish world in cam frame
                                SE3 T_cam_from_world(frame->T_f_w_* T_world_from_vision_);
                                q = Quaterniond(T_cam_from_world.rotation_matrix());
                                p = T_cam_from_world.translation();
                                Cov = frame->Cov_;
                        }
                        else
                        {
                                // publish cam in world frame
                                SE3 T_world_from_cam(T_world_from_vision_*frame->T_f_w_.inverse());
                                q = Quaterniond(T_world_from_cam.rotation_matrix()*T_world_from_vision_.rotation_matrix().transpose());
                                p = T_world_from_cam.translation();
                                Cov = T_world_from_cam.Adj()*frame->Cov_*T_world_from_cam.inverse().Adj();
                        }
                        geometry_msgs::PoseWithCovarianceStampedPtr msg_pose(new geometry_msgs::PoseWithCovarianceStamped);
                        msg_pose->header = header_msg;
                        msg_pose->pose.pose.position.x = p[0];
                        msg_pose->pose.pose.position.y = p[1];
                        msg_pose->pose.pose.position.z = p[2];
                        msg_pose->pose.pose.orientation.x = q.x();
                        msg_pose->pose.pose.orientation.y = q.y();
                        msg_pose->pose.pose.orientation.z = q.z();
                        msg_pose->pose.pose.orientation.w = q.w();
                        for(size_t i=0; i<36; ++i)
                                msg_pose->pose.covariance[i] = Cov(i%6, i/6);
                        pub_pose_.publish(msg_pose);    // #pub 3
                }
        }

        void Visualizer::visualizeMarkers(
                        const FramePtr& frame,
                        const set<FramePtr>& core_kfs,
                        const Map& map)
        {
                if(frame == NULL)
                        return;

                vk::output_helper::publishTfTransform(
                                frame->T_f_w_*T_world_from_vision_.inverse(),
                                ros::Time(frame->timestamp_), "cam_pos", "world", br_);

                if(pub_frames_.getNumSubscribers() > 0 || pub_points_.getNumSubscribers() > 0)
                {
                        vk::output_helper::publishHexacopterMarker(
                                        pub_frames_, "cam_pos", "cams", ros::Time(frame->timestamp_),
                                        1, 0, 0.3, Vector3d(0.,0.,1.)); // #pub 4

                        Vector3d traj_color;
                        traj_color << frame->align_method_/255.0, 0, 0.5;
//                        switch(frame->align_method_)
//                        {
//                                case 1:
//                                        traj_color<<0,0.5,0;
//                                        break;
//                                case 2:
//                                        traj_color<<0.5,0,0;
//                                        break;
//                                default:
//                                        traj_color<<0,0,0.5;
//                                        break;
//                        }
                        vk::output_helper::publishPointMarker(
                                        pub_points_, T_world_from_vision_*frame->pos(), "trajectory",
                                        ros::Time::now(), trace_id_, 0, 0.012, traj_color);    // #pub 5
                        if(frame->isKeyframe() || publish_map_every_frame_)
                                publishMapRegion(core_kfs);
                        removeDeletedPts(map);
                }
        }

        void Visualizer::visualizeDenseCloud(
                        FramePtr frame, bool draw_uncertainty)
        {
                list<Vector3d> cloud;
                visualization_msgs::Marker line_list;
                if (frame)
                {
                        boost::unique_lock<boost::mutex> lock(frame->depth_map_mut_);
                        if (frame->depth_map_.empty()) {
                                return;
                        }

                        for(auto it=frame->depth_map_.begin(), ite=frame->depth_map_.end(); it != ite; ++it)
                        {
                                if (!it->converged) {
                                        continue;
                                }
                                cloud.push_back(T_world_from_vision_*frame->T_f_w_.inverse()*(1.0/it->mu * it->ftr->f)); 
                                if (!draw_uncertainty) {
                                        continue;
                                }
                                geometry_msgs::Point p;
                                Vector3d line_start = T_world_from_vision_*frame->T_f_w_.inverse()
                                        *(1.0/(it->mu + sqrt(it->sigma2)) * it->ftr->f);
                                Vector3d line_end = T_world_from_vision_*frame->T_f_w_.inverse()
                                        *(1.0/max(it->mu - sqrt(it->sigma2), 0.00000001f) * it->ftr->f);
                                p.x = line_start(0);
                                p.y = line_start(1);
                                p.z = line_start(2);
                                line_list.points.push_back(p);
                                p.x = line_end(0);
                                p.y = line_end(1);
                                p.z = line_end(2);
                                line_list.points.push_back(p);
                        }
                        vk::output_helper::publishPointCloud(pub_points_, cloud, "cloud",
                                        ros::Time::now(), frame->id_, 0, 0.008, 
                                        Vector3d(0.0, 0.3, 0.0), publish_points_display_time_);
                        if (draw_uncertainty) {
                                vk::output_helper::publishLineGroup(pub_points_, line_list, 
                                                "uncertainty",
                                                ros::Time::now(), frame->id_, 0, 0.008, 
                                                Vector3d(0.0, 0.3, 0.0), publish_points_display_time_);
                        }
                }
        }


        void Visualizer::publishMapRegion(set<FramePtr> frames)
        {
                if(pub_points_.getNumSubscribers() > 0)
                {
                        int ts = vk::Timer::getCurrentTime();
                        for(set<FramePtr>::iterator it=frames.begin(); it!=frames.end(); ++it)
                                displayKeyframeWithMps(*it, ts);
                }
        }

        void Visualizer::removeDeletedPts(const Map& map)
        {
                if(pub_points_.getNumSubscribers() > 0)
                {
                        for(list<Point*>::const_iterator it=map.trash_points_.begin(); it!=map.trash_points_.end(); ++it)
                                vk::output_helper::publishPointMarker(pub_points_, Vector3d(), "pts", ros::Time::now(), (*it)->id_, 2, 0.006, Vector3d());
                }
        }

        void Visualizer::displayKeyframeWithMps(const FramePtr& frame, int ts)
        {
                // publish keyframe
                SE3 T_world_cam(T_world_from_vision_*frame->T_f_w_.inverse());
                vk::output_helper::publishFrameMarker(
                                pub_frames_, T_world_cam.rotation_matrix(),
                                T_world_cam.translation(), "kfs", ros::Time::now(), frame->id_*10, 0, 0.015);

                // publish point cloud and links
                for(Features::iterator it=frame->fts_.begin(); it!=frame->fts_.end(); ++it)
                {
                        if((*it)->point == NULL)
                                continue;

                        if((*it)->point->last_published_ts_ == ts)
                                continue;

                        vk::output_helper::publishPointMarker(
                                        pub_points_, T_world_from_vision_*(*it)->point->pos_, "pts",
                                        ros::Time::now(), (*it)->point->id_, 0, 0.005, Vector3d(1.0, 0., 1.0),
                                        publish_points_display_time_);
                        (*it)->point->last_published_ts_ = ts;
                }
        }

        void Visualizer::exportToDense(const FramePtr& frame)
        {
                // publish air_ground_msgs
                if(frame != NULL && dense_pub_nth_ > 0
                                && trace_id_%dense_pub_nth_ == 0 && pub_dense_.getNumSubscribers() > 0)
                {
                        svo_msgs::DenseInput msg;
                        msg.header.stamp = ros::Time(frame->timestamp_);
                        msg.header.frame_id = "/world";
                        msg.frame_id = frame->id_;

                        cv_bridge::CvImage img_msg;
                        img_msg.header.stamp = msg.header.stamp;
                        img_msg.header.frame_id = "camera";
                        img_msg.image = frame->img();
                        img_msg.encoding = sensor_msgs::image_encodings::MONO8;
                        msg.image = *img_msg.toImageMsg();

                        double min_z = std::numeric_limits<double>::max();
                        double max_z = std::numeric_limits<double>::min();
                        for(Features::iterator it=frame->fts_.begin(); it!=frame->fts_.end(); ++it)
                        {
                                if((*it)->point==NULL)
                                        continue;
                                Vector3d pos = frame->T_f_w_*(*it)->point->pos_;
                                min_z = fmin(pos[2], min_z);
                                max_z = fmax(pos[2], max_z);
                        }
                        msg.min_depth = (float) min_z;
                        msg.max_depth = (float) max_z;

                        // publish cam in world frame
                        SE3 T_world_from_cam(T_world_from_vision_*frame->T_f_w_.inverse());
                        Quaterniond q(T_world_from_cam.rotation_matrix());
                        Vector3d p(T_world_from_cam.translation());

                        msg.pose.position.x = p[0];
                        msg.pose.position.y = p[1];
                        msg.pose.position.z = p[2];
                        msg.pose.orientation.w = q.w();
                        msg.pose.orientation.x = q.x();
                        msg.pose.orientation.y = q.y();
                        msg.pose.orientation.z = q.z();
                        pub_dense_.publish(msg);
                }
        }

} // end namespace svo
